package kudu

import org.apache.kudu.client.CreateTableOptions
import org.apache.kudu.spark.kudu.KuduContext
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types.{IntegerType, LongType, StringType, StructField, StructType}
import org.junit.{Before, Test}

/**
 * Spark 操作kudu demo
 */
class KuduSparkAPI {

  @Test
  def createTable(): Unit = {
    // 1、创建KuduContext 和SparkSession
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("sparkOpKudu")
      .config("spark.some.config.option", "some-value")
      .getOrCreate()

    val KUDU_MASTER = "s202:7051,s203:7051,s204:7051"
    val kuduContext = new KuduContext(KUDU_MASTER, spark.sparkContext)


    // 判断表是否存在
    val tableName = "sparkopKudu"
    if (kuduContext.tableExists(tableName)) kuduContext.deleteTable(tableName)


    // 创建表
    val schema = StructType(
      StructField("id", LongType, nullable = false) ::
        StructField("name", StringType, nullable = false) ::
        StructField("age", IntegerType, nullable = false) :: Nil
    )
    val keys = Seq("id")

    import scala.collection.JavaConverters._
    val options = new CreateTableOptions()
      .setRangePartitionColumns(List("id").asJava)
      .setNumReplicas(1)

    kuduContext.createTable(tableName,
      schema,
      keys,
      options
    )
  }


  @Test
  def curd(): Unit = {
    // 1、创建KuduContext 和SparkSession
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("sparkOpKudu")
      .config("spark.some.config.option", "some-value")
      .getOrCreate()

    // 2、创建kuduContext
    val KUDU_MASTER = "s202:7051,s203:7051,s204:7051"
    val kuduContext = new KuduContext(KUDU_MASTER, spark.sparkContext)

    import spark.implicits._
    val df = Seq(Person(1, "ling", 18), Person(2, "chb", 28)).toDF()
    val tableName = "sparkopKudu"

    // 删： 只需要key
    kuduContext.deleteRows(df.select("id"), tableName)

    // 增
    kuduContext.insertRows(df, tableName)

    // 改
    kuduContext.updateRows(df, tableName)
  }
}

// case class Person(id: Long, name: String, age: Int)
